import glob

import torch
from PIL import Image

import numpy as np
from torch.utils.data import random_split, DataLoader

import matplotlib
import matplotlib.pyplot as plt

from AnimaiClassModel import MonkeyClassificationModel

img_dir = 'H:\\workspace\\pythons\\data\\'


label_dict = {
               0: "alouattapalliata",
               1: "erythrocebuspatas",
               2: "cacajaocalvus",
               3: "macacafuscata",
               4: "cebuellapygmea",
               5: "cebuscapucinus",
               6: "micoargentatus",
               7: "saimirisciureus",
               8: "aotusnigriceps",
               9: "trachypithecusjohnii"
}

def getImage(path):
    im = Image.open(path)
    im = im.resize((400, 300))
    pixels = np.asarray(im).astype('float32')
    pixels /= 255.0
    pixels = torch.from_numpy(pixels)
    return pixels

def predict_image(input_img, model):
    inputs = input_img.unsqueeze(0)
    predictions = model(inputs)
    _, preds  = torch.max(predictions, dim=1)
    return preds[0].item()

def real_run():
    model = torch.load(img_dir + 'model\\monkey_classification.pth')

    img = getImage(img_dir + "validation\\validation\\n6\\n6013.jpg")
    # print(img)

    pred = predict_image(img, model)
    print(label_dict[pred])



# real_run()